import torch
import torch.nn as nn
from transformers import BertModel
from config import Config


class BertNLU(nn.Module):
    def __init__(self, num_intents, num_slots):
        super().__init__()
        # 加载预训练BERT模型
        self.bert = BertModel.from_pretrained(Config.bert_model)
        # 意图分类头（输入维度=BERT隐藏层维度）
        self.intent_head = nn.Linear(Config.bert_hidden_size, num_intents)
        # 槽位分类头（输入维度=BERT隐藏层维度）
        self.slot_head = nn.Linear(Config.bert_hidden_size, num_slots)
        # Dropout层（防止过拟合）
        self.dropout = nn.Dropout(Config.dropout)

    def forward(self, input_ids, attention_mask):
        """前向传播：输入文本→意图和槽位预测"""
        # BERT输出：(batch_size, seq_len, hidden_size)
        bert_out = self.bert(
            input_ids=input_ids,
            attention_mask=attention_mask
        ).last_hidden_state  # 取最后一层隐藏状态

        # 意图预测：使用[CLS]位置的输出（batch_size, hidden_size）
        intent_logits = self.intent_head(self.dropout(bert_out[:, 0, :]))

        # 槽位预测：使用全序列输出（batch_size, seq_len, hidden_size）
        slot_logits = self.slot_head(self.dropout(bert_out))

        return intent_logits, slot_logits
